import logging
import os
import re
from pathlib import Path

from tqdm import tqdm

from gitmanager.projectstate import GitProject
from llm.client import DeepSeekClient, APIMode
from prompts.prompts import Prompt
from utils import system_call, load_jsonl, read_file, rename_file_to_temp, write_file, rename_file_from_temp, save_jsonl


class RustFixer:
    def __init__(self, git_rust_project: GitProject,
                 client: DeepSeekClient,
                 prompt: Prompt,
                 log_path,
                 rounds=5,
                 progress=False):
        self.log_path = log_path
        self.git_rust_project = git_rust_project

        self.rust_repo_path = Path(self.git_rust_project.source_path).resolve()

        self.client = client
        self.prompt = prompt

        self.rounds = rounds
        self.progress = progress
        self.progress_bar = None

    def cargo_check_code(self):
        current_dir = os.getcwd()
        os.chdir(self.rust_repo_path)
        return_code, stdout, stderr = system_call('cargo check --message-format json')
        os.chdir(current_dir)
        if return_code == 0:
            return []
        return self.parse_check_messages(stdout)

    def parse_check_messages(self, json_output):
        output = []
        messages = load_jsonl(json_output)
        try:
            for message in messages[:-1]:
                if 'message' not in message:
                    continue
                parsed = self.parse_message(message)
                if parsed is None:
                    continue
                # todo：目前只关注error
                if parsed['level'] == 'warning':
                    continue
                output.append(parsed)
        except Exception as e:
            logger = logging.getLogger()
            jsonl_file = self.log_path / f'cargo_error.jsonl'
            logger.error(f"解析cargo check错误 {e}: | 错误文件存到 {jsonl_file}")
            save_jsonl(jsonl_file, messages)
        return output

    @staticmethod
    def parse_message(message):
        parsed_message = {}
        message = message['message']
        if message['level'] not in {'warning', 'error'}:
            return None
        parsed_message['level'] = message['level']
        parsed_message['rendered'] = message['rendered']

        spans = message['spans']
        primary_span = None
        for span in spans:
            if span['is_primary']:
                primary_span = span
                break

        if primary_span is None:
            return None

        parsed_message['file_name'] = primary_span['file_name']
        parsed_message['line_start'] = primary_span['line_start']
        parsed_message['line_end'] = primary_span['line_end']
        parsed_message['column_start'] = primary_span['column_start']
        parsed_message['column_end'] = primary_span['column_end']
        parsed_message['label'] = primary_span['label']
        parsed_message['suggested_replacement'] = primary_span['suggested_replacement']
        parsed_message['suggestion_applicability'] = primary_span['suggestion_applicability']
        return parsed_message

    def fix_rust_code(self, rust_file=None):
        if rust_file is None:
            get_feed_back = lambda : list(self.cargo_check_code())
        else:
            rust_file = rust_file.relative_to(self.rust_repo_path)
            get_feed_back = lambda : list(filter(lambda x:
                                                 Path(x['file_name']) == Path(rust_file),
                                                 self.cargo_check_code()))

        feedback = get_feed_back()

        if self.progress:
            self.progress_bar = tqdm(total=len(feedback), desc='Fixing')

        i = 0
        rounds = self.rounds
        while len(feedback) > 0 and rounds > 0 and i < len(feedback):
            error_num = len(feedback)
            error = feedback[i]
            error_code_path = self.rust_repo_path / error['file_name']
            fixed_rust_code = self.fix_one_error(error)
            temp_file = self.apply_patch(fixed_rust_code, error_code_path)
            if temp_file is None:  # 没有正常修复
                return False
            new_feedback = get_feed_back()
            if len(new_feedback) < error_num:
                if self.progress:
                    self.progress_bar.n = len(new_feedback)
                    self.progress_bar.refresh()

                # self.git_rust_project.commit([error_code_path], error['rendered'])

                if len(new_feedback) == 0:
                    return True
                os.remove(temp_file)
                feedback = new_feedback
                i = 0
                rounds = self.rounds
            else:
                rename_file_from_temp(temp_file)
                i = (i + 1) % len(feedback)
                rounds -= 1

        jsonl_file = self.log_path / f'fix.jsonl'
        self.client.save_histories(jsonl_file)

        if len(feedback) > 0:
            feedback = get_feed_back()

        if len(feedback) == 0:
            return True
        return False

    def fix_one_error(self, error):
        error_code_path = self.rust_repo_path / error['file_name']
        error_code = read_file(error_code_path)
        actual_prompt = self.prompt(error=error['rendered'], error_code=error_code)
        response = self.client.generate(actual_prompt, mode=APIMode.COMPLETION)
        return response

    @staticmethod
    def apply_patch(fixed_rust_code, rust_code_path):
        m = re.search(r"```rust\n(.*?)```", fixed_rust_code, re.DOTALL)
        if m:
            rust_code = m.group(1)
            temp_file = rename_file_to_temp(rust_code_path)
            write_file(rust_code_path, rust_code)
            return temp_file
        else:
            logger = logging.getLogger()
            logger.error(f'模型返回的响应没有代码: {rust_code_path}')
            return None

    def __del__(self):
        if self.progress_bar is not None:
            self.progress_bar.close()
